#include "OpenGLShader.h"

#include "OpenGLBase.h"

#include <fstream>

namespace SE
{
	static inline GLenum ToGLenum(const String& type)
    {
        if (type == "vertex")
            return GL_VERTEX_SHADER;
        if (type == "fragment" || type == "pixel")
            return GL_FRAGMENT_SHADER;

        SE_ASSERT(false, "Unknown shader type!");

        return GL_NONE;
    }

	static inline String ToString(const GLenum stage)
    {
        switch (stage)
        {
            default: SE_ASSERT(false, "Unknown shader stage!"); return "";
            case GL_VERTEX_SHADER:   return "GL_VERTEX_SHADER";
            case GL_FRAGMENT_SHADER: return "GL_FRAGMENT_SHADER";
        }
    }

	OpenGLShader::OpenGLShader(const String& filePath) : mFilePath(filePath)
	{
		SE_PROFILE_FUNCTION();

		String source = readFile(filePath);
		auto sources = preProcess(source);
		compile(sources);

		// Extract name from filePath
		auto lastSlash = filePath.find_last_of("/\\");
		lastSlash = lastSlash == String::npos ? 0 : lastSlash + 1;
		auto lastDot = filePath.rfind('.');
		auto count = lastDot == String::npos ? filePath.size() - lastSlash : lastDot - lastSlash;
		mName = filePath.substr(lastSlash, count);
	}

	OpenGLShader::OpenGLShader(const String& name, const String& vertexSrc, const String& fragmentSrc) : mName(name)
	{
		SE_PROFILE_FUNCTION();

		HashMap<U32, String> sources;
		sources[GL_VERTEX_SHADER]   = vertexSrc;
		sources[GL_FRAGMENT_SHADER] = fragmentSrc;
		compile(sources);
	}

	OpenGLShader::~OpenGLShader()
	{
		SE_PROFILE_FUNCTION();

		glDeleteProgram(mRenderID);
	}

	String OpenGLShader::readFile(const String& filePath)
	{
		SE_PROFILE_FUNCTION();

		String result;
		std::ifstream in(filePath, std::ios::in | std::ios::binary); // ifstream closes itself due to RAII.
		if (in)
		{
			in.seekg(0, std::ios::end);
			U64 size = in.tellg();
			if (size != -1)
			{
				result.resize(size);
				in.seekg(0, std::ios::beg);
				in.read(&result[0], size);
			}
			else
			{
				LOG_ERROR("Could not read from file '{0}'.", filePath);
			}
		}
		else
		{
			LOG_ERROR("Could not open file '{0}'.", filePath);
		}

		return result;
	}

	HashMap<U32, String> OpenGLShader::preProcess(const String& source)
	{
		SE_PROFILE_FUNCTION();

		HashMap<U32, String> shaderSources;

		const char* typeToken = "#type";
		U64 typeTokenLength = strlen(typeToken);
		U64 pos = source.find(typeToken, 0); //Start of shader type declaration line
		while (pos != String::npos)
		{
			U64 eol = source.find_first_of("\r\n", pos); //End of shader type declaration line
			SE_ASSERT(eol != String::npos, "Syntax error");
			U64 begin = pos + typeTokenLength + 1; //Start of shader type name (after "#type " keyword)
			String type = source.substr(begin, eol - begin);
			SE_ASSERT(ToGLenum(type), "Invalid shader type specified");

			U64 nextLinePos = source.find_first_not_of("\r\n", eol); //Start of shader code after shader type declaration line
			SE_ASSERT(nextLinePos != String::npos, "Syntax error");
			pos = source.find(typeToken, nextLinePos); //Start of next shader type declaration line

			shaderSources[(U32)ToGLenum(type)] = (pos == String::npos) ? source.substr(nextLinePos) : source.substr(nextLinePos, pos - nextLinePos);
		}

		return shaderSources;
	}

	void OpenGLShader::compile(const HashMap<U32, String>& shaderSources)
	{
		SE_PROFILE_FUNCTION();

		GLuint program = glCreateProgram();
		SE_ASSERT(shaderSources.size() <= 2, "We only support 2 shaders for now.");
		Array<U32, 2> glShaderIDs;
		int glShaderIDIndex = 0;
		for (auto& kv : shaderSources)
		{
			GLenum type = kv.first;
			const String& source = kv.second;

			GLuint shader = glCreateShader(type);

			const GLchar* sourceCStr = source.c_str();
			glShaderSource(shader, 1, &sourceCStr, 0);

			glCompileShader(shader);

			GLint isCompiled = 0;
			glGetShaderiv(shader, GL_COMPILE_STATUS, &isCompiled);
			if (isCompiled == GL_FALSE)
			{
				GLint maxLength = 0;
				glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &maxLength);

				Vector<GLchar> infoLog(maxLength);
				glGetShaderInfoLog(shader, maxLength, &maxLength, &infoLog[0]);

				glDeleteShader(shader);

				LOG_ERROR("{0}", infoLog.data());
				SE_ASSERT(false, "Shader compilation failed!");
				break;
			}

			glAttachShader(program, shader);
			glShaderIDs[glShaderIDIndex++] = shader;
		}
		
		mRenderID = program;

		// Link our program
		glLinkProgram(program);

		// Note the different functions here: glGetProgram* instead of glGetShader*.
		GLint isLinked = 0;
		glGetProgramiv(program, GL_LINK_STATUS, (int*)&isLinked);
		if (isLinked == GL_FALSE)
		{
			GLint maxLength = 0;
			glGetProgramiv(program, GL_INFO_LOG_LENGTH, &maxLength);

			// The maxLength includes the NULL character
			Vector<GLchar> infoLog(maxLength);
			glGetProgramInfoLog(program, maxLength, &maxLength, &infoLog[0]);

			// We don't need the program anymore.
			glDeleteProgram(program);
			
			for (auto id : glShaderIDs)
				glDeleteShader(id);

			LOG_ERROR("{0}", infoLog.data());
			SE_ASSERT(false, "Shader link failed!");
			return;
		}

		for (auto id : glShaderIDs)
		{
			glDetachShader(program, id);
			glDeleteShader(id);
		}
	}

	void OpenGLShader::bind() const
	{
		SE_PROFILE_FUNCTION();

		glUseProgram(mRenderID);
	}

	void OpenGLShader::unbind() const
	{
		SE_PROFILE_FUNCTION();

		glUseProgram(0);
	}

	void OpenGLShader::setUniformInt(const String& name, const I32 value)
	{
		SE_PROFILE_FUNCTION();

		glUniform1i(getUniformLocation(name), value);
	}

	void OpenGLShader::setUniformIntArray(const String& name, I32* values, const U32 count)
	{
		SE_PROFILE_FUNCTION();

		glUniform1iv(getUniformLocation(name), count, values);
	}

	void OpenGLShader::setUniformFloat(const String& name, const Float value)
	{
		SE_PROFILE_FUNCTION();

		glUniform1f(getUniformLocation(name), value);
	}

	void OpenGLShader::setUniformFloat2(const String& name, const Vector2f& value)
	{
		SE_PROFILE_FUNCTION();

		glUniform2f(getUniformLocation(name), value.x, value.y);
	}

	void OpenGLShader::setUniformFloat3(const String& name, const Vector3f& value)
	{
		SE_PROFILE_FUNCTION();

		glUniform3f(getUniformLocation(name), value.x, value.y, value.z);
	}

	void OpenGLShader::setUniformFloat4(const String& name, const Vector4f& value)
	{
		SE_PROFILE_FUNCTION();

		glUniform4f(getUniformLocation(name), value.x, value.y, value.z, value.w);
	}

	void OpenGLShader::setUniformMatrix3f(const String& name, const Matrix3f& matrix)
	{
		SE_PROFILE_FUNCTION();

		glUniformMatrix3fv(getUniformLocation(name), 1, GL_FALSE, glm::value_ptr(matrix));
	}

	void OpenGLShader::setUniformMatrix4f(const String& name, const Matrix4f& matrix)
	{
		SE_PROFILE_FUNCTION();

		glUniformMatrix4fv(getUniformLocation(name), 1, GL_FALSE, glm::value_ptr(matrix));
	}

    const String &OpenGLShader::getName() const
    {
        return mName;
    }

    I32 OpenGLShader::getUniformLocation(const String& name)
    {
		SE_PROFILE_FUNCTION();

		if (mUniformLocationCache.find(name) != mUniformLocationCache.end())
			return mUniformLocationCache[name];

		I32 location = glGetUniformLocation(mRenderID, name.c_str());
		if (location == -1)
			LOG_WARN("Uniform '{}' does not exist !", name);
		mUniformLocationCache[name] = location;

		return location;
    }
} // !namespace SE